"""
测试流域边界获取功能
基于 mghydro 项目的流域划分功能，通过指定出口点坐标来获取准确的流域边界

参考: https://github.com/mheberger/delineator
"""

import os
import sys
import traceback

import pandas as pd
import geopandas as gpd
from shapely.geometry import Point
import tempfile
import shutil
from loguru import logger

# 添加项目根目录到路径，以便导入模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from delineate import delineate
from config import *


class WatershedDelineator:
    """
    流域边界获取器
    通过出口点坐标获取准确的流域边界
    """
    
    def __init__(self, data_dir=None, output_dir=None):
        """
        初始化流域边界获取器
        
        Args:
            data_dir: 数据目录路径，如果为None则使用config.py中的默认设置
            output_dir: 输出目录路径，如果为None则使用临时目录
        """
        self.data_dir = data_dir
        self.output_dir = output_dir or tempfile.mkdtemp()
        
    def create_outlet_csv(self, outlets, csv_path):
        """
        创建出口点CSV文件
        
        Args:
            outlets: 出口点列表，每个元素为字典，包含id, lat, lng, name等字段
            csv_path: CSV文件保存路径
        """
        df = pd.DataFrame(outlets)
        
        # 确保包含必需的列
        required_columns = ['id', 'lat', 'lng', 'name']
        for col in required_columns:
            if col not in df.columns:
                if col == 'id':
                    df['id'] = range(1, len(df) + 1)
                elif col == 'name':
                    df['name'] = [f'Outlet_{i}' for i in range(1, len(df) + 1)]
                else:
                    raise ValueError(f"缺少必需的列: {col}")
        
        df.to_csv(csv_path, index=False)
        return csv_path
    
    def delineate_watershed(self, lat, lng, outlet_name="Custom_Outlet", outlet_id=1, 
                          high_res=True, area_estimate=None):
        """
        通过单个出口点获取流域边界
        
        Args:
            lat: 纬度 (decimal degrees)
            lng: 经度 (decimal degrees) 
            outlet_name: 出口点名称
            outlet_id: 出口点ID
            high_res: 是否使用高分辨率模式
            area_estimate: 预估流域面积 (km²)，可选
            
        Returns:
            GeoDataFrame: 包含流域边界的地理数据框
        """
        
        # 创建出口点数据
        outlet_data = {
            'id': outlet_id,
            'lat': lat,
            'lng': lng,
            'name': outlet_name
        }
        
        if area_estimate:
            outlet_data['area'] = area_estimate
            
        return self.delineate_watersheds([outlet_data], high_res=high_res)
    
    def delineate_watersheds(self, outlets, high_res=True):
        """
        通过多个出口点获取流域边界
        
        Args:
            outlets: 出口点列表，每个元素为字典
            high_res: 是否使用高分辨率模式
            
        Returns:
            GeoDataFrame: 包含所有流域边界的地理数据框
        """
        
        # 创建临时CSV文件
        temp_csv = os.path.join(self.output_dir, 'temp_outlets.csv')
        self.create_outlet_csv(outlets, temp_csv)
        
        # 备份原始配置
        original_outlets_csv = OUTLETS_CSV
        original_high_res = HIGH_RES
        original_output_dir = OUTPUT_DIR
        original_output_ext = OUTPUT_EXT
        
        try:
            # 临时修改全局配置
            import config
            config.OUTLETS_CSV = temp_csv
            config.HIGH_RES = high_res
            config.OUTPUT_DIR = self.output_dir
            config.OUTPUT_EXT = "gpkg"  # 使用GeoPackage格式
            
            # 执行流域划分
            delineate()
            
            # 读取结果
            output_files = []
            for file in os.listdir(self.output_dir):
                if file.endswith('.gpkg') and 'watershed' in file.lower():
                    output_files.append(os.path.join(self.output_dir, file))
            
            if not output_files:
                raise FileNotFoundError("未找到输出的流域边界文件")
            
            # 合并所有结果
            all_watersheds = []
            for file in output_files:
                gdf = gpd.read_file(file)
                all_watersheds.append(gdf)
            
            if len(all_watersheds) == 1:
                return all_watersheds[0]
            else:
                return gpd.GeoDataFrame(pd.concat(all_watersheds, ignore_index=True))
        except Exception as e:
            traceback.print_exc()
        finally:
            # 恢复原始配置
            config.OUTLETS_CSV = original_outlets_csv
            config.HIGH_RES = original_high_res
            config.OUTPUT_DIR = original_output_dir
            config.OUTPUT_EXT = original_output_ext
            
            # 清理临时文件
            if os.path.exists(temp_csv):
                os.remove(temp_csv)
    
    def save_watershed(self, watershed_gdf, output_path, format='geojson'):
        """
        保存流域边界到文件
        
        Args:
            watershed_gdf: 流域边界GeoDataFrame
            output_path: 输出文件路径
            format: 输出格式 ('geojson', 'gpkg', 'shp')
        """
        if format.lower() == 'geojson':
            watershed_gdf.to_file(output_path, driver='GeoJSON')
        elif format.lower() == 'gpkg':
            watershed_gdf.to_file(output_path, driver='GPKG')
        elif format.lower() == 'shp':
            watershed_gdf.to_file(output_path)
        else:
            raise ValueError(f"不支持的格式: {format}")
    
    def cleanup(self):
        """清理临时文件"""
        if self.output_dir and os.path.exists(self.output_dir):
            shutil.rmtree(self.output_dir)


def test_single_outlet():
    """测试单个出口点的流域划分"""
    logger.info("测试单个出口点流域划分...")
    
    delineator = WatershedDelineator()
    
    try:
        # 冰岛 Nordhura River at Stekkur 测试点
        outlet_id = "6401070"
        lat = 64.71072
        lng = -21.60337
        outlet_name = "Nordhura River at Stekkur"
        area_estimate = 507  # 预估面积 507 km²
        
        logger.info(f"正在为出口点 ({lat}, {lng}) 划分流域...")
        logger.info(f"出口点ID: {outlet_id}, 名称: {outlet_name}")
        
        # 获取流域边界
        watershed = delineator.delineate_watershed(
            lat=lat, 
            lng=lng, 
            outlet_name=outlet_name,
            outlet_id=outlet_id,
            high_res=True,
            area_estimate=area_estimate
        )
        
        if watershed is not None and len(watershed) > 0:
            logger.success(f"成功获取流域边界!")
            logger.info(f"流域数量: {len(watershed)}")
            logger.info(f"流域总面积: {watershed.geometry.area.sum():.2f} 平方度")
            
            # 保存结果为 GeoJSON
            output_path = os.path.join(delineator.output_dir, "test_single_watershed.geojson")
            delineator.save_watershed(watershed, output_path, format='geojson')
            logger.success(f"流域边界已保存到: {output_path}")
            
            return True
        else:
            logger.error(f"测试失败: 未能获取流域边界")
            return False
            
    except Exception as e:
        logger.error(f"测试失败: {str(e)}")
        return False
    finally:
        delineator.cleanup()


def test_multiple_outlets():
    """测试多个出口点的流域划分"""
    logger.info("测试多个出口点流域划分...")
    
    delineator = WatershedDelineator()
    
    try:
        # 定义多个测试出口点（冰岛 Nordhura River at Stekkur 及其附近点）
        outlets = [
            {
                'id': '6401070',
                'lat': 64.71072,
                'lng': -21.60337,
                'name': 'Nordhura River at Stekkur',
                'area': 507
            },
            {
                'id': '6401071',
                'lat': 64.71000,
                'lng': -21.60000,
                'name': 'Nordhura River nearby point',
                'area': 400
            }
        ]
        
        logger.info(f"正在为 {len(outlets)} 个出口点划分流域...")
        
        # 获取多个流域边界
        watersheds = delineator.delineate_watersheds(outlets, high_res=True)
        
        if watersheds is not None and len(watersheds) > 0:
            logger.success(f"成功获取流域边界!")
            logger.info(f"流域数量: {len(watersheds)}")
            
            # 保存结果为 GeoJSON
            output_path = os.path.join(delineator.output_dir, "test_multiple_watersheds.geojson")
            delineator.save_watershed(watersheds, output_path, format='geojson')
            logger.success(f"流域边界已保存到: {output_path}")
            
            return True
        else:
            logger.error(f"测试失败: 未能获取流域边界")
            return False
            
    except Exception as e:
        logger.error(f"测试失败: {str(e)}")
        return False
    finally:
        delineator.cleanup()


if __name__ == "__main__":
    # 检查必要的数据目录
    logger.info("=== 流域边界获取测试 ===")
    logger.info("基于 mghydro 项目的流域划分功能")
    logger.info("参考: https://github.com/mheberger/delineator")
    
    # 检查数据目录
    required_dirs = [LOWRES_CATCHMENTS_DIR, HIGHRES_CATCHMENTS_DIR]
    missing_dirs = [d for d in required_dirs if not os.path.exists(d)]
    
    if missing_dirs:
        logger.warning("以下数据目录不存在，请先下载相应的数据文件:")
        for d in missing_dirs:
            logger.warning(f"  - {d}")
        logger.info("请参考 README.md 中的数据下载说明")
        logger.info("或访问: https://mghydro.com/watersheds/")
    else:
        # 运行测试
        logger.info("开始运行测试...")
        test_single_outlet()
        test_multiple_outlets()
        
        logger.info("=== 使用说明 ===")
        logger.info("1. 创建 WatershedDelineator 实例")
        logger.info("2. 调用 delineate_watershed() 方法获取单个流域")
        logger.info("3. 调用 delineate_watersheds() 方法获取多个流域")
        logger.info("4. 使用 save_watershed() 方法保存结果")
        logger.info("示例代码:")
        logger.info("```python")
        logger.info("delineator = WatershedDelineator()")
        logger.info("watershed = delineator.delineate_watershed(lat=64.71072, lng=-21.60337)")
        logger.info("delineator.save_watershed(watershed, 'my_watershed.geojson', format='geojson')")
        logger.info("```")